"""
角色扮演Web界面
提供角色选择和语音对话的用户界面，支持流式文本和语音输出
"""

import streamlit as st
import requests
import base64
import json
import time
import sseclient
from typing import Dict, Any, List

# API基础URL
API_BASE_URL = "http://localhost:8000/role-play"

def get_all_roles() -> List[Dict[str, Any]]:
    """获取所有角色"""
    response = requests.get(f"{API_BASE_URL}/roles")
    if response.status_code == 200:
        return response.json()
    return []

def select_role(role_id: str) -> Dict[str, Any]:
    """选择角色"""
    response = requests.get(f"{API_BASE_URL}/select/{role_id}")
    if response.status_code == 200:
        return response.json()
    return None

# 暂时注释单轮对话功能，改用流式对话
# def chat_with_role(message: str, role_id: str = None) -> Dict[str, Any]:
#     """与角色对话"""
#     response = requests.post(
#         f"{API_BASE_URL}/chat", 
#         json={"message": message, "role_id": role_id}
#     )
#     if response.status_code == 200:
#         return response.json()
#     return {"success": False, "error": "API请求失败"}

def chat_with_role_stream(message: str, role_id: str = None):
    """与角色进行流式对话，支持流式文本和语音输出"""
    headers = {'Accept': 'text/event-stream'}
    response = requests.post(
        f"{API_BASE_URL}/chat/stream",
        json={"message": message, "role_id": role_id},
        headers=headers,
        stream=True
    )
    
    if response.status_code != 200:
        yield {"success": False, "error": "API请求失败"}
        return
    
    client = sseclient.SSEClient(response)
    for event in client.events():
        yield json.loads(event.data)

def play_audio(audio_base64: str):
    """播放Base64编码的音频"""
    if not audio_base64:
        return  # 如果音频数据为空，直接返回
    try:
        audio_bytes = base64.b64decode(audio_base64)
        st.audio(audio_bytes, format="audio/wav")
    except Exception as e:
        st.error(f"音频解码错误: {str(e)}")

def main():
    st.set_page_config(
        page_title="AI角色扮演",
        page_icon="🎭",
        layout="wide"
    )
    
    st.title("🎭 AI角色扮演")
    st.subheader("与你喜欢的角色进行语音聊天")
    
    # 侧边栏 - 角色选择
    with st.sidebar:
        st.header("角色选择")
        
        # 显示角色列表
        if "roles" not in st.session_state:
            # 初始加载所有角色
            st.session_state.roles = get_all_roles()
        
        st.subheader("可用角色")
        
        for role in st.session_state.roles:
            if st.button(f"选择 {role['name']}", key=f"select_{role['id']}"):
                st.session_state.current_role = role
                st.session_state.messages = []  # 清空对话历史
                st.rerun()
            st.caption(role['description'])
    
    # 主界面 - 聊天区域
    if "current_role" in st.session_state:
        role = st.session_state.current_role
        
        st.header(f"与 {role['name']} 对话")
        st.caption(role['description'])
        
        # 初始化消息历史
        if "messages" not in st.session_state:
            st.session_state.messages = []
        
        # 显示消息历史
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.write(message["content"])
                if "audio" in message:
                    play_audio(message["audio"])
        
        # 用户输入
        user_input = st.chat_input("输入你的消息...")
        
        if user_input:
            # 显示用户消息
            with st.chat_message("user"):
                st.write(user_input)
            
            # 添加到历史
            st.session_state.messages.append({"role": "user", "content": user_input})
            
            # 获取角色流式回复
            with st.chat_message("assistant"):
                message_placeholder = st.empty()
                audio_placeholder = st.empty()
                full_response = ""
                audio_chunks = []
                
                # 创建一个空的消息占位
                st.session_state.messages.append({
                    "role": "assistant", 
                    "content": "",
                    "audio": None
                })
                
                # 使用流式API获取回复
                for chunk in chat_with_role_stream(user_input, role_id=role["id"]):
                    # 检查chunk是否为字符串（可能是原始文本响应）
                    if isinstance(chunk, str):
                        full_response = chunk
                        message_placeholder.markdown(full_response)
                        st.session_state.messages[-1]["content"] = full_response
                        continue
                        
                    # 检查chunk是否为字典且包含success字段
                    if isinstance(chunk, dict) and "success" in chunk and not chunk["success"]:
                        st.error(chunk.get("error", "未知错误"))
                        break
                    
                    # 更新文本
                    if isinstance(chunk, dict) and "text" in chunk:
                        full_response = chunk["text"]
                        message_placeholder.markdown(full_response)
                        # 更新会话历史
                        st.session_state.messages[-1]["content"] = full_response
                    
                    # 收集音频块
                    if isinstance(chunk, dict) and "audio" in chunk and chunk["audio"]:
                        audio_chunks.append(chunk["audio"])
                        # 播放最新的音频块
                        try:
                            # 直接使用二进制数据播放音频
                            audio_data = base64.b64decode(chunk["audio"])
                            audio_placeholder.audio(audio_data, format="audio/wav")
                        except Exception as e:
                            st.error(f"音频播放错误: {str(e)}")
                    
                    # 如果是最终响应
                    if isinstance(chunk, dict) and chunk.get("is_final", False):
                        break
                
                # 如果收集了音频块，将它们合并并存储
                if audio_chunks:
                    st.session_state.messages[-1]["audio"] = audio_chunks[-1]  # 使用最后一个音频块作为完整音频
    else:
        st.info("👈 请从侧边栏选择一个角色开始对话")

if __name__ == "__main__":
    main()